# 333
import torch
import torch.nn as nn
from timm.models.layers import DropPath, to_2tuple, trunc_normal_
import torch.utils.checkpoint as checkpoint
import numpy as np
import torch.nn.functional as F
import copy
import math
from models.CaptionModel import CaptionModel
from torch.nn.utils.rnn import PackedSequence, pack_padded_sequence, pad_packed_sequence
from torch.autograd import Variable


class Feature_fusion(CaptionModel):
    def __init__(self, opt):
        super(Feature_fusion, self).__init__()
        self.opt = opt
        self.vocab_size = opt.vocab_size
        self.seq_length = opt.seq_length
        self.img_size = opt.image_size
        self.patch_size = opt.patch_size
        self.in_chans = opt.in_chans
        self.depths = opt.swin_depths
        # self.num_heads = opt.swin_num_heads
        # self.embed_dim = opt.embed_dim
        # self.window_size = opt.window_size
        self.att_feat_size = opt.att_feat_size
        self.feature_fuse_size = opt.feature_fuse_size
        # self.mlp_ratio = opt.mlp_ratio
        # self.qkv_bias = opt.qkv_bias
        # self.qk_scal = opt.qk_scale
        # self.drop_rate = opt.drop_rate
        # self.attn_drop_rate = 0.
        self.drop_path_rate = opt.drop_path_rate
        self.drop_prob_lm = opt.drop_prob_lm
        self.ape = opt.ape
        # self.num_layers = len(opt.swin_depths)
        self.patch_norm = opt.patch_norm
        # self.num_features = int(self.embed_dim * 2 ** (self.num_layers - 1))
        self.norm_layer = nn.LayerNorm
        # self.use_checkpoint = opt.use_checkpoint

        self.patch_embed = PatchEmbed(img_size=self.img_size, patch_size=16, in_chans=self.in_chans,
                                      embed_dim=self.opt.patch_embed_dim)
        self.position_embed = PositionalEncoding(self.opt.patch_embed_dim, self.opt.drop_rate)
        self.lin_lay = nn.Linear(in_features=self.opt.patch_embed_dim, out_features=self.opt.input_encoding_size, bias=True)
        tgt_vocab = self.vocab_size + 1

        # Relation Transformer model
        self.att_embed = nn.Sequential(*(
                (nn.Linear(self.att_feat_size, self.opt.input_encoding_size),
                 nn.ReLU(),
                 nn.Dropout(self.drop_prob_lm))))

        self.relation_model = self.make_model(
            0, tgt_vocab, N=self.opt.encode_num_layers, d_model=self.opt.input_encoding_size, d_ff=self.opt.rnn_size)

        self.patch_relation_model = self.make_model(
            0, tgt_vocab, N=8, h=8, d_model=self.opt.patch_embed_dim, d_ff=self.opt.rnn_size)

        # Bulid decoder
        c = copy.deepcopy
        position = PositionalEncoding(self.opt.input_encoding_size, self.opt.drop_rate)

        self.tgt_embed = nn.Sequential(Embeddings(self.opt.input_encoding_size, tgt_vocab), c(position))
        self.decoder = self.make_decoder(0, tgt_vocab, N=self.opt.decoder_num_layers,
                                         d_model=self.opt.input_encoding_size, d_ff=self.opt.rnn_size)
        self.generator = Generator(self.opt.input_encoding_size, tgt_vocab)

        self._init_param([self.patch_embed, self.position_embed, self.lin_lay, self.att_embed, self.patch_relation_model,self.relation_model, position, self.tgt_embed,
                          self.decoder, self.generator])

        # self.apply(self._init_weights)

    def _init_param(self, input):
        for model in input:
            for p in model.parameters():
                if p.dim() > 1:
                    nn.init.xavier_uniform_(p)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    def make_decoder(self, src_vocab, tgt_vocab, N=6, d_model=512, d_ff=2048, h=8, drop_out=0.1):
        c = copy.deepcopy

        attn = MultiHeadedAttention(h, d_model)
        ff = PositionwiseFeedForward(d_model, d_ff, drop_out)

        model = Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), drop_out), N)

        return model

    def make_model(self, src_vocab, tgt_vocab, N=6, d_model=512, d_ff=2048, h=8, drop_out=.1):
        c = copy.deepcopy

        bbox_attn = MultiHeadedAttention(h, d_model)

        ff = PositionwiseFeedForward(d_model, d_ff, drop_out)

        model = Encoder(EncoderLayer(d_model, c(bbox_attn), c(ff), drop_out), N)

        return model

    def _prepare_feature(self, att_feats, att_masks=None, seq=None):
        att_feats, att_masks = self.clip_att(att_feats, att_masks)
        att_feats = pack_wrapper(self.att_embed, att_feats, att_masks)

        if att_masks is None:
            att_masks = att_feats.new_ones(att_feats.shape[:2], dtype=torch.long)
        att_masks = att_masks.unsqueeze(-2)

        if seq is not None:
            seq = seq[:, :-1]
            seq_mask = (seq.data > 0)
            seq_mask[:, 0] = 1

            # new_seq_mask = torch.Tensor(seq_mask.size(), dtype=torch.uint8)
            # for i in range(seq_mask.size(0)):
            #     for j in range(seq_mask.size(1)):
            #         if seq_mask[i, j] == True:
            #             new_seq_mask[i, j] = 1
            #         else:
            #             new_seq_mask[i, j] = 0
            # seq_mask = new_seq_mask
            seq_mask = seq_mask.unsqueeze(-2)

            seq_mask = seq_mask & subsequent_mask(seq.size(-1)).to(seq_mask)
        else:
            seq_mask = None
        return att_feats, seq, att_masks, seq_mask

    def decode(self, encoder_out, seq, seq_mask, att_masks):
        return self.decoder(self.tgt_embed(seq), encoder_out, seq_mask, att_masks)

    def _forward(self, img, fc_feats, att_feats, seq, masks, att_masks=None):

        imgs = self.patch_embed(img)
        imgs = self.position_embed(imgs)
        patch_masks = None
        patch_attn = self.patch_relation_model(imgs, patch_masks)
        patch_attn = self.lin_lay(patch_attn)

        # bbox_attn
        att_feats, seq, att_masks, seq_mask = self._prepare_feature(att_feats, att_masks, seq)

        bbox_attn, att_masks = self.set_att(att_feats, patch_attn, att_masks)

        bbox_attn = self.relation_model(bbox_attn, att_masks)

        # decode
        decoder_out = self.decode(bbox_attn, seq, seq_mask, att_masks)
        outputs = self.generator(decoder_out)
        return outputs

    def set_att(self, att_feats, encoder_out, att_masks):
        """
        This function will concat the att_feats and encoder_out
        :param att_feats: output of relation transformer
        :param encoder_out: out put of swin transformer
        :param att_masks: masks of att_feats
        :return: att_feats ＆ att_masks
        """
        batch_size = att_feats.size(0)
        att_masks = att_masks.squeeze(1)

        att_feats_size = att_feats.size()
        patch_size = encoder_out.size()
        att_masks_size = att_masks.size()
        att_masks_num = att_masks.data.long().sum(1)
        new_att_feats = torch.zeros(att_feats_size[0], att_masks_size[1]+patch_size[1], att_feats_size[2], dtype=torch.float32).cuda()
        for i in range(att_masks_size[0]):
            new_att_feats[i, :att_masks_num[i], :] = att_feats[i, :att_masks_num[i], :]
            new_att_feats[i, att_masks_num[i]: att_masks_num[i]+patch_size[1], :] = encoder_out[i, :, :]
            new_att_feats[i, att_masks_num[i]+patch_size[1]:, :] = att_feats[i, att_masks_num[i]:, :]
        new_att_masks = torch.zeros(att_masks_size[0], att_masks_size[1]+patch_size[1], dtype=torch.int64).cuda()
        att_masks_num += patch_size[1]
        for j in range(att_masks_size[0]):
            new_att_masks[j, :att_masks_num[j]] = 1
        new_att_masks = new_att_masks.unsqueeze(1)

        return new_att_feats, new_att_masks

    def _sample(self, img, fc_feats, att_feats, att_masks=None, opt={}):
        sample_max = opt.get('sample_max', 1)
        beam_size = opt.get('beam_size', 1)
        temperature = opt.get('temperature', 1.0)
        decoding_constraint = opt.get('decoding_constraint', 0)
        if beam_size > 1:
            return self._sample_beam(img, fc_feats, att_feats, att_masks, opt)

        batch_size = img.shape[0]

        state = None

        imgs = self.patch_embed(img)
        imgs = self.position_embed(imgs)
        patch_masks = None
        patch_attn = self.patch_relation_model(imgs, patch_masks)
        patch_attn = self.lin_lay(patch_attn)

        att_feats, seq, att_masks, seq_mask = self._prepare_feature(att_feats, att_masks)

        # att_feats, encoder_out 使用特征融合方法进行特征融合

        bbox_attn, att_masks = self.set_att(att_feats, patch_attn, att_masks)

        bbox_attn = self.relation_model(bbox_attn, att_masks)

        seq = att_feats.new_zeros((batch_size, self.seq_length), dtype=torch.long)
        seqLogprobs = att_feats.new_zeros(batch_size, self.seq_length)

        for t in range(self.seq_length + 1):
            if t == 0:  # input <bos>
                it = fc_feats.new_zeros(batch_size, dtype=torch.long)

            logprobs, state = self.get_logprobs_state(it, bbox_attn, att_masks, state)

            # sample the next word
            if t == self.seq_length:
                break
            if sample_max:
                sampleLogprobs, it = torch.max(logprobs.data, 1)
                it = it.view(-1).long()
            else:
                if temperature == 1.0:
                    prob_prev = torch.exp(logprobs.data)  # fetch prev distribution: shape Nx(M+1)
                else:
                    # scale logprobs by temperature
                    prob_prev = torch.exp(torch.div(logprobs.data, temperature))
                it = torch.multinomial(prob_prev, 1)
                sampleLogprobs = logprobs.gather(1, it)  # gather the logprobs at sampled positions
                it = it.view(-1).long()  # and flatten indices for downstream processing

            if t == 0:
                unfinished = it > 0
            else:
                unfinished = unfinished * (it > 0)
            it = it * unfinished.type_as(it)
            seq[:, t] = it
            seqLogprobs[:, t] = sampleLogprobs.view(-1)

            if unfinished.sum() == 0:
                break

        return seq, seqLogprobs

    def _sample_beam(self, img, fc_feats, att_feats, att_masks=None, opt={}):
        beam_size = opt.get('beam_size', 5)
        batch_size = fc_feats.size(0)

        imgs = self.patch_embed(img)
        imgs = self.position_embed(imgs)
        patch_masks = None
        patch_attn = self.patch_relation_model(imgs, patch_masks)
        patch_attn = self.lin_lay(patch_attn)

        att_feats, seq, att_masks, seq_mask = self._prepare_feature(att_feats, att_masks)

        bbox_attn, att_masks = self.set_att(att_feats, patch_attn, att_masks)

        bbox_attn = self.relation_model(bbox_attn, att_masks)

        assert beam_size <= self.vocab_size + 1, 'lets assume this for now, otherwise this corner case causes a few headaches down the road. can be dealt with in future if needed'
        seq = torch.LongTensor(self.seq_length, batch_size).zero_()
        seqLogprobs = torch.FloatTensor(self.seq_length, batch_size)
        # lets process every image independently for now, for simplicity

        self.done_beams = [[] for _ in range(batch_size)]
        for k in range(batch_size):
            state = None
            tmp_memory = bbox_attn[k:k + 1].expand(*((beam_size,) + bbox_attn.size()[1:])).contiguous().cuda()
            tmp_att_masks = att_masks[k:k + 1].expand(
                *((beam_size,) + att_masks.size()[1:])).contiguous().cuda() if att_masks is not None else None

            for t in range(1):
                if t == 0:  # input <bos>
                    it = fc_feats.new_zeros([beam_size], dtype=torch.long)

                logprobs, state = self.get_logprobs_state(it, tmp_memory, tmp_att_masks, state)

            self.done_beams[k] = self.beam_search(state, logprobs, tmp_memory, tmp_att_masks, opt=opt)
            seq[:, k] = self.done_beams[k][0]['seq']  # the first beam has highest cumulative score
            seqLogprobs[:, k] = self.done_beams[k][0]['logps']
        # return the samples and their log likelihoods
        return seq.transpose(0, 1), seqLogprobs.transpose(0, 1)

    def get_logprobs_state(self, it, memory, mask, state):
        if state is None:
            ys = it.unsqueeze(1)
        else:
            ys = torch.cat([state[0][0], it.unsqueeze(1)], dim=1)
        out = self.decode(memory, ys, subsequent_mask(ys.size(1)).to(memory.device), mask)
        logprobs = self.generator(out[:, -1])

        return logprobs, [ys.unsqueeze(0)]

    def clip_att(self, att_feats, att_masks):
        # Clip the length of att_masks and att_feats to the maximum length
        if att_masks is not None:
            max_len = att_masks.data.long().sum(1).max()
            att_feats = att_feats[:, :max_len].contiguous()
            att_masks = att_masks[:, :max_len].contiguous()
        return att_feats, att_masks


class PatchEmbed(nn.Module):
    """ Image to Patch Embedding
    """

    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
        super().__init__()
        img_size = (img_size, img_size)
        patch_size = (patch_size, patch_size)
        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
        self.img_size = img_size
        self.patch_size = patch_size
        self.num_patches = num_patches

        self.project = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)

    def forward(self, x):
        B, C, H, W = x.shape
        # FIXME look at relaxing size constraints
        assert H == self.img_size[0] and W == self.img_size[1], \
            f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
        x = self.project(x)
        x = x.flatten(2)
        x = x.transpose(1, 2)
        return x


class Decoder(nn.Module):
    """Generic N layer decoder with masking."""

    def __init__(self, layer, N):
        super(Decoder, self).__init__()
        self.layers = clones(layer, N)
        self.norm = LayerNorm(layer.size)

    def forward(self, x, encoder_out, tgt_mask, att_masks):
        for layer in self.layers:
            x = layer(x, encoder_out, tgt_mask, att_masks)
        return self.norm(x)


class DecoderLayer(nn.Module):
    """Decoder is mode of self-attn, src-attn, and feed forward(defined below)"""

    def __init__(self, size, self_attn, src_attn, feed_forward, dropout):
        super(DecoderLayer, self).__init__()
        self.size = size
        self.self_attn = self_attn
        self.src_attn = src_attn
        self.feed_forward = feed_forward
        self.sub_layer = clones(SublayerConnection(size, dropout), 3)

    def forward(self, x, encoder_out, tgt_mask, att_masks):
        m = encoder_out
        x = self.sub_layer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask))
        x = self.sub_layer[1](x, lambda x: self.src_attn(x, m, m, att_masks))
        return self.sub_layer[2](x, self.feed_forward)


class SublayerConnection(nn.Module):
    """
    A residual connection followed by a layer norm.
    Note for code simplicity the norm is first as opposed to last.
    """

    def __init__(self, size, dropout):
        super(SublayerConnection, self).__init__()
        self.norm = LayerNorm(size)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, sublayer):
        """Apply residual connection to any sublayer with the same size"""
        return x + self.dropout(sublayer(self.norm(x)))


class LayerNorm(nn.Module):
    """Construct a layernorm module (See citation for details)"""

    def __init__(self, features, eps=1e-6):
        super(LayerNorm, self).__init__()
        self.a_2 = nn.Parameter(torch.ones(features))
        self.b_2 = nn.Parameter(torch.zeros(features))
        self.eps = eps

    def forward(self, x):
        mean = x.mean(-1, keepdim=True)
        std = x.std(-1, keepdim=True)
        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2


class Embeddings(nn.Module):
    def __init__(self, d_model, vocab):
        super(Embeddings, self).__init__()
        self.lut = nn.Embedding(vocab, d_model)
        self.d_model = d_model

    def forward(self, x):
        return self.lut(x) * math.sqrt(self.d_model)


class PositionalEncoding(nn.Module):
    """Implement the PE function"""

    def __init__(self, d_model, dropout, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)

        # Compute the positional encodings once in log space.
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1).float()
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe)

    def forward(self, x):
        x = x + Variable(self.pe[:, :x.size(1)], requires_grad=False)
        return self.dropout(x)


class MultiHeadedAttention(nn.Module):
    def __init__(self, h, d_model, dropout=0.1):
        """Take in model size and number of heads."""
        super(MultiHeadedAttention, self).__init__()
        assert d_model % h == 0

        self.d_k = d_model // h
        self.h = h
        self.linears = clones(nn.Linear(d_model, d_model), 4)
        self.attn = None
        self.dropout = nn.Dropout(p=dropout)

    def forward(self, query, key, value, mask=None):
        if mask is not None:
            mask = mask.unsqueeze(1)
        # else:
        #     key = key.unsqueeze(1)
        #     value = value.unsqueeze(1)
        nbatches = query.size(0)
        # 1)
        query, key, value = [l(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2) for l, x in
                             zip(self.linears, (query, key, value))]
        # 2)
        x, self.attn = attention(query, key, value, mask=mask, dropout=self.dropout)
        # 3)
        x = x.transpose(1, 2).contiguous().view(nbatches, -1, self.h * self.d_k)
        return self.linears[-1](x)


class PositionwiseFeedForward(nn.Module):
    """Implements FFN equation."""

    def __init__(self, d_model, d_ff, dropout=0.1):
        super(PositionwiseFeedForward, self).__init__()
        self.w_1 = nn.Linear(d_model, d_ff)
        self.w_2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        return self.w_2(self.dropout(F.relu(self.w_1(x))))


class Generator(nn.Module):
    """Define standard linear + softmax generation step."""

    def __init__(self, d_model, vocab):
        super(Generator, self).__init__()
        self.proj = nn.Linear(d_model, vocab)

    def forward(self, x):
        return F.log_softmax(self.proj(x), dim=-1)


class Encoder(nn.Module):
    def __init__(self, layer, N):
        super(Encoder, self).__init__()
        self.layers = clones(layer, N)
        self.norm = LayerNorm(layer.size)

    def forward(self, x, mask):
        for layer in self.layers:
            x = layer(x, mask)
        return self.norm(x)


class EncoderLayer(nn.Module):
    def __init__(self, size, self_attn, feed_forward, drop_out):
        super(EncoderLayer, self).__init__()
        self.self_attn = self_attn
        self.feed_forward = feed_forward
        self.sublayer = clones(SublayerConnection(size, drop_out), 2)
        self.size = size

    def forward(self, x, mask):
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))
        return self.sublayer[1](x, self.feed_forward)


def window_partition(x, window_size):
    """
    :param x:(B, H, W, C)
    :param window_size:(int) window_size
    :return:
        windows:(num_windows*B, window_size, window_size, C)
    """
    B, H, W, C = x.shape
    x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
    windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
    return windows


def window_reverse(windows, window_size, H, W):
    """
    :param windows: (num_windows*B, window_size, window_size, C)
    :param window_size: window size
    :param H: Height of images
    :param W: width of images
    :return: x:(B, H, W, C)
    """
    B = int(windows.shape[0] / (H * W / window_size / window_size))
    x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
    x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
    return x


def clones(module, N):
    """Produce N identical layers."""
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])


def attention(query, key, value, mask=None, dropout=None):
    """Compute Scaled dot Product Attention"""
    d_k = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    p_attn = F.softmax(scores, dim=-1)
    if dropout is not None:
        p_attn = dropout(p_attn)
    return torch.matmul(p_attn, value), p_attn


def subsequent_mask(size):
    """Mask out subsequent positions."""
    attn_shape = (1, size, size)
    subsequent_mask = np.triu(np.ones(attn_shape), k=1).astype('uint8')
    return torch.from_numpy(subsequent_mask) == 0


def pack_wrapper(module, att_feats, att_masks):
    if att_masks is not None:
        packed, inv_ix = sort_pack_padded_sequence(att_feats, att_masks.data.long().sum(1))
        return pad_unsort_packed_sequence(PackedSequence(module(packed[0]), packed[1]), inv_ix)
    else:
        return module(att_feats)


def sort_pack_padded_sequence(input, lengths):
    sorted_lengths, indices = torch.sort(lengths, descending=True)
    tmp = pack_padded_sequence(input[indices], sorted_lengths.cpu(), batch_first=True)
    inv_ix = indices.clone()
    inv_ix[indices] = torch.arange(0, len(indices)).type_as(inv_ix)
    return tmp, inv_ix


def pad_unsort_packed_sequence(input, inv_ix):
    tmp, _ = pad_packed_sequence(input, batch_first=True)
    tmp = tmp[inv_ix]
    return tmp


